import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim

# 模拟数据
# 假设有3个类别
num_classes = 2
# 假设有一个输入特征张量，batch_size = 5
# inputs = torch.randn(5, requires_grad=True)
inputs = torch.tensor([[0, 0, 1], [1, 1, 1], [1, 0, 1], [0, 1, 1]]).float()
# 假设目标标签是一个包含类别索引的张量
targets = torch.tensor([0, 1, 1, 0])

# 定义模型（定义一个全连接层）
linear = nn.Linear(in_features=3, out_features=num_classes)
# 定义优化器
optimizer = optim.SGD(linear.parameters(), lr=0.01)  # 学习率设置为0.01
# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    # 假设我们每轮迭代100次
    for _ in range(10):
        # 前向传播（进行线性变换）
        logits = linear(inputs)
        # 计算损失
        loss = criterion(F.log_softmax(logits, dim=1), targets)
        # 反向传播和优化
        optimizer.zero_grad()  # 清空梯度
        loss.backward()  # 反向传播
        optimizer.step()  # 更新权重

for _ in  range(10):
    # 测试模型
    linear.eval()  # 将模型设置为评估模式
    with torch.no_grad():  # 禁用梯度计算
        outputs = linear(torch.tensor([0, 1, 0]).float())
        #print(f"outputs:{outputs}")
        # 维度索引是从0开始的，-1通常表示最后一个维度，所有如果outputs是一维张量，那么dim=0
        predicted = outputs.argmax(dim=0, keepdim=True)
        print(f"predicted:{predicted}")
        # correct = predicted.eq(targets.view_as(predicted)).sum().item()
        # accuracy = correct / targets.size(0)
        # print(f"Test Accuracy: {accuracy:.4f}")

if __name__ == '__main__':
    print("end")
